import os
import json
import subprocess
import logging
from typing import Dict, List, Any, Optional
from pathlib import Path
import re

from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate

from openai import OpenAI
from dotenv import  load_dotenv
import textwrap
import argparse
# 加载环境变量
load_dotenv()

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class CodeGenerator:
    def __init__(self, api_key: str, mapping_rules_path: str,base_url:str,model_name:str):
        """
        初始化代码生成器
        
        Args:
            qwen_api_key: Qwen API 密钥
            mapping_rules_path: 映射规则 JSON 文件路径
            base_url:api base url
            model_name:model name
        """
        self.api_key = api_key
        self.mapping_rules = self._load_mapping_rules(mapping_rules_path)
        self.base_url = base_url 
        self.model_name = model_name 
        self.llm = self._setup_llm()
        self.function_mapper = {
        "read_specify_files": self.read_specify_files
        }

    def _load_mapping_rules(self, mapping_rules_path: str) -> Dict:
        """加载映射规则"""
        try:
            with open(mapping_rules_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            logger.error(f"加载映射规则失败: {e}")
            raise
    
    def _setup_llm(self):
        """设置 Qwen LLM"""
        # 注意：这里使用 OpenAI 兼容的接口，实际使用时需要根据 Qwen API 调整
        return OpenAI(
            api_key=self.api_key,
            base_url=self.base_url,
        )
    
    def format_call(self,user_content:list,format_type:str="json_object"):
        response = self.llm.chat.completions.create(
            model=self.model_name,  # 根据实际情况调整
            messages=user_content,
            response_format={"type":format_type}
        )
        #logger.info(f"content:{user_content}")
        json_str = response.choices[0].message.content
        #logger.info(f"json的内容:{json_str} ")
        result = json.loads(json_str)
        # 提取header和code字段（带默认值，避免字段缺失）
        header = result.get("header", [])
        code = result.get("code", "")
        return header,code

    def stream_call(self,user_content:list,tools:list=None):
        response = self.llm.chat.completions.create(
            model=self.model_name,  # 根据实际情况调整
            messages=user_content,
            stream=True,
            tools=tools,
            parallel_tool_calls=True,
            extra_body={"enable_thinking": True, "thinking_budget": 1000}
        )
        return response

    def process_response(self,response):
        reasoning_content = ""  # 完整思考过程
        answer_content = ""  # 完整回复
        tool_info = []  #工具调用信息
        is_answering = False  # 是否进入回复阶段
        print("\n" + "=" * 20 + "思考过程" + "=" * 20 + "\n")

        for chunk in response:
            if not chunk.choices:
                print("\nUsage:")
                print(chunk.usage)
                continue

            delta = chunk.choices[0].delta

         # 只收集思考内容
            if hasattr(delta, "reasoning_content") and delta.reasoning_content is not None:
                if not is_answering:
                    print(delta.reasoning_content, end="", flush=True)
                reasoning_content += delta.reasoning_content

        # 收到content，开始进行回复
            if hasattr(delta, "content") and delta.content:
                if not is_answering:
                    print("\n" + "=" * 20 + "完整回复" + "=" * 20 + "\n")
                    is_answering = True
                print(delta.content, end="", flush=True)
                answer_content += delta.content

             # 处理工具调用信息（支持并行工具调用）
            if delta.tool_calls is not None:
                for tool_call in delta.tool_calls:
                    index = tool_call.index  # 工具调用索引，用于并行调用
                    
                    # 动态扩展工具信息存储列表
                    while len(tool_info) <= index:
                        tool_info.append({})
                    
                    # 收集工具调用ID（用于后续函数调用）
                    if tool_call.id:
                        tool_info[index]['id'] = tool_info[index].get('id', '') + tool_call.id
                    
                    # 收集函数名称（用于后续路由到具体函数）
                    if tool_call.function and tool_call.function.name:
                        tool_info[index]['name'] = tool_info[index].get('name', '') + tool_call.function.name
                    
                    # 收集函数参数（JSON字符串格式，需要后续解析）
                    if tool_call.function and tool_call.function.arguments:
                        tool_info[index]['arguments'] = tool_info[index].get('arguments', '') + tool_call.function.arguments
        res=[reasoning_content,answer_content,tool_info]
        return res

    def read_specify_files(self,path:str,name:str)->str:
        target_names = []
        if name.endswith(('.h', '.hpp','.d.ts')):
            target_names.append(name)
        else:
            target_names.append(f"{name}.h")
            target_names.append(f"{name}.hpp")
            target_names.append(f"{name}.d.ts")
        content=""
        # 搜索文件
        merged_content = []
        found_files = []
        search_path=Path(path)
        for ext in ['*.h', '*.hpp','*.d.ts']:
            for file in search_path.rglob(ext):
                if file.name in target_names:
                    found_files.append(file)
                    with open(file, 'r', encoding='utf-8') as f:
                        content = f.read()
                    merged_content.append(content)
                    logger.info(f"找到头文件: {file}")
        
        if not found_files:
            logger.info(f"在 {search_path} 中未找到 {name} 相关头文件")
            return ""
        result = '\n'.join(merged_content)
        return result


    def read_header_files(self, header_dir: str) -> Dict[str, str]:
        """
        读取头文件内容
        
        Args:
            header_dir: 头文件目录路径
            
        Returns:
            字典：文件名 -> 文件内容
        """
        header_files = {}
        header_dir_path = Path(header_dir)
        
        # 查找所有 .h 和 .hpp 文件 .d.ts文件
        for ext in ['*.h', '*.hpp','*.d.ts']:
            for header_file in header_dir_path.glob(ext):
                try:
                    with open(header_file, 'r', encoding='utf-8') as f:
                        content = f.read()
                    header_files[header_file.name] = content
                    logger.info(f"成功读取头文件: {header_file.name}")
                except Exception as e:
                    logger.error(f"读取头文件 {header_file} 失败: {e}")
        
        return header_files
    
    def tools_prompt(self)->list:
        tools=[
        {
            "type": "function",
            "function": {
                "name": "read_specify_files",
                "description": "当前代码中有未定义类型时调用来查询头文件内容",
                "parameters": {
                    "type":"object",
                    "properties":{
                        "name":{
                            "type":"string",
                            "description":"要查询的头文件名称"
                        }
                    },
                    "required":["name"]
                }  
            }
        }
        ]
        
        return tools  

    def process_content(self,content:str,filename:str,project_path:str)->str:
        """
        调用API处理头文件内容,过滤私有成员
        
        Args:
            content: 原始头文件内容
            filename:文件名
        Returns:
            过滤私有成员后的头文件内容
        """
        prompt_template=textwrap.dedent(
        """
        请对输入的内容进行处理,提取其中的头文件和代码内容
        处理规则：
        1.头文件提取非标准库头文件
        2.代码内容只保留各种类型声明和其内部Public成员、方法的声明签名,以及全局function的声明签名,删除其他内容,特别是private内容一定要删除
        3.总结提取后用到的语言本身未提供的自定义类型,如果有记得调用"read_specify_files"函数
        
        输出格式：
        '''
        header:用实际的header替换
        code:用处理后的code替换
        undefined:用未定义类型替换
        """) 
        prompt_search=textwrap.dedent(
        """
        请根据提供的补充文件,查询补充文件中是否有目标文件中未定义类型的定义,如果有将其定义加入到code中并删除undefined中的该类型
        参考文件：
        {file_content}
        输出格式：
        '''
        header:原目标文件的header
        code:在原目标文件的code基础上添加找到的未定义类型的定义
        undefined:原目标文件的未定义类型删去已找到定义的
        """)
        input_template= textwrap.dedent(
        """
        文件名：
        {filename}
        文件内容：
        {content}
        """
        )
        tools=self.tools_prompt()
        prompt=PromptTemplate.from_template(prompt_template).format()
        input=PromptTemplate.from_template(input_template).format(filename=filename,content=content)
        messages = [{"role": "system", "content": prompt},{"role":"user","content":input}]
        try:
            code=""
            count=2
            ans=[]
            while(count):
                response = self.stream_call(messages,tools)
                #获取完整的内容:思考，回复，工具调用
                ans = self.process_response(response)
                logger.info(ans[1])
                if not ans[2]:
                    logger.info("无工具调用")
                    code=ans[1]
                    break
                else:
                    logger.info("工具调用")
                    delta_all_content=[]
                    for call in ans[2]:
                        func=self.function_mapper[call["name"]]
                        args_json=call["arguments"]
                        args=json.loads(args_json)
                        delta_content=func(project_path,**args)
                        delta_all_content.append(delta_content)
                    all_content="\n".join(delta_all_content)
                    search_prompt=PromptTemplate.from_template(prompt_search).format(file_content=all_content)
                    #messages = [{"role":"system","content":prompt},{"role":"assistant","content":ans[1]},{"role":"user","content":search_prompt}]
                    messages.append({"role":"assistant","content":ans[1]})
                    messages.append({"role":"user","content":search_prompt})
                count-=1    

            logger.info(f"成功处理{filename}的内容 ")

            return code
        
        except json.JSONDecodeError:
            raise
        except Exception as e:
            logger.error(f"生成代码失败: {e}")
            raise
        
    
    def generate_code_prompt(self, header_content: str, filename: str) -> str:
        """
        生成代码生成提示
        
        Args:
            header_content: 头文件内容
            filename: 头文件名
            
        Returns:
            格式化后的提示词
        """
        prompt_template =textwrap.dedent(
        """
        请根据以下头文件内容和映射规则生成对应的idl代码,
        
        映射规则:
        {mapping_rules}
        
        头文件名称: 
        {filename}
        头文件内容:
        ```
        {header_content}
        ```
        
        生成的代码应有良好的书写规范,处理好换行与缩进,只返回代码内容，不要包含任何解释或注释。
        生成的代码应该能够直接编译运行。
        生成的代码:
        ```taihe
        """)
        prompt=PromptTemplate.from_template(prompt_template).format(
            mapping_rules=json.dumps(self.mapping_rules, indent=2, ensure_ascii=False),
            filename=filename,
            header_content=header_content
            )
        messages = [{"role": "user", "content": prompt}]
        return messages
        
    
    def generate_repair_prompt(self, original_header: str, generated_code: str, 
                              error_message: str, filename: str) -> str:
        """
        生成修复代码的提示
        
        Args:
            original_header: 原始头文件内容
            generated_code: 生成的代码
            error_message: 错误信息
            filename: 文件名
            
        Returns:
            格式化后的修复提示词
        """
        prompt_template =textwrap.dedent(
        """
        你生成的代码编译或测试失败，请根据错误信息进行修复。
        
        原始头文件内容:
        ```cpp
        {original_header}
        ```
        
        你之前生成的代码:
        ```taihe
        {generated_code}
        ```
        
        错误信息:
        {error_message}
        
        映射规则：
        {mapping_rules}
        
        请修复代码中的问题，只返回修复后的完整代码，不要包含任何解释。
        
        修复后的代码:
        ```taihe
        """)
        
        return PromptTemplate.from_template(prompt_template).format(
            original_header=original_header,
            generated_code=generated_code,
            error_message=error_message,
            mapping_rules=json.dumps(self.mapping_rules, indent=2, ensure_ascii=False),
        )
    
    def extract_code_from_response(self, response: str) -> str:
        """
        从模型响应中提取代码
        
        Args:
            response: 模型响应文本
            
        Returns:
            提取的代码内容
        """
        
        # 尝试提取代码块
        code_pattern = r'```(?:taihe|cpp)?\s*(.*?)\s*```'
        matches = re.findall(code_pattern, response, re.DOTALL)
        
        if matches:
            return matches[0].strip()
        
        # 如果没有代码块，返回整个响应'''
        return response
    
    def generate_code(self, header_content: str, filename: str) -> str:
        """
        生成代码
        
        Args:
            header_content: 头文件内容
            filename: 文件名
            
        Returns:
            生成的代码
        """
        prompt = self.generate_code_prompt(header_content, filename)
        try:
            logger.info("开始生成代码")
            response= self.stream_call(prompt)
            ans = self.process_response(response)
            code = self.extract_code_from_response(ans[1])
            logger.info(f"成功为 {filename} 生成代码")
            return code
        except Exception as e:
            logger.error(f"生成代码失败: {e}")
            raise
    
    def repair_code(self, original_header: str, generated_code: str, 
                   error_message: str, filename: str) -> str:
        """
        修复代码
        
        Args:
            original_header: 原始头文件
            generated_code: 生成的代码
            error_message: 错误信息
            filename: 文件名
            
        Returns:
            修复后的代码
        """
        prompt = self.generate_repair_prompt(original_header, generated_code, 
                                           error_message, filename)
        message=[{"role":"user","content":prompt}]
        try:
            chat = self.stream_call(message)
            ans = self.process_response(chat) 
            repaired_code = self.extract_code_from_response(ans[1])
            logger.info(f"成功修复 {filename} 的代码")
            return repaired_code
        except Exception as e:
            logger.error(f"修复代码失败: {e}")
            raise

    def save_generated_code(self, code: str, output_path: str, filename: str):
        """
        保存生成的代码
        
        Args:
            code: 代码内容
            output_path: 输出目录
            filename: 文件名
        """
        save_name = filename
        if filename.endswith('.d.ts'):
            save_name  = filename[:-5]
        else:
            save_name   = Path(filename).stem
        output_file = Path(output_path) / f"{Path(save_name)}.taihe"
        
        try:
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(code)
            logger.info(f"代码已保存到: {output_file}")
        except Exception as e:
            logger.error(f"保存代码失败: {e}")
            raise

        return output_file

    def run_taihec_script(self,script_path: str, code_file: str,output_path:str) -> tuple:
        """
        运行 taihec 脚本检验
        
        Args:
            script_path: 脚本路径
            code_file: 代码文件路径
            # 完整命令拆分：每个部分作为列表的一个元素
            # 命令解析：
            #       . D:/taihe/bin/taihec.bat       -> 脚本路径（注意开头的.是source命令，Windows下通常可省略）
            # ./generate/Recorder.taihe      -> 输入文件
            # -O ./generate                  -> 输出目录选项
            # -G napi-bridge cpp-author      -> 生成器选项及参数
        Returns:
            (成功与否, 错误信息)
        """
        try:
            command = [
                script_path,  # 脚本路径
                code_file,
                "-O", output_path,
                "-G", "napi-bridge", "cpp-author",
            ]
            result = subprocess.run(
                command,
                capture_output=True,
                text=True,
                timeout=300  # 5分钟超时
            )
            
            if result.returncode == 0:
                logger.info("脚本检验通过")
                return True, ""
            else:
                error_msg = f"脚本检验失败:\nSTDOUT: {result.stdout}\nSTDERR: {result.stderr}"
                logger.error(error_msg)
                return False, error_msg
                
        except subprocess.TimeoutExpired:
            error_msg = "脚本执行超时"
            logger.error(error_msg)
            return False, error_msg
        except Exception as e:
            error_msg = f"执行脚本异常: {e}"
            logger.error(error_msg)
            return False, error_msg
    
    def search_taihe_dir(self,script_path:str,target:str):
        try:
            command = [script_path, target]
            result = subprocess.run(
                command,
                capture_output=True,
                text=True,
                timeout=300  # 5分钟超时
            )
            
            if result.returncode == 0:

                return result.stdout.strip()
            else:
                error_msg = f"cmake脚本检验失败:\nSTDOUT: {result.stdout}\nSTDERR: {result.stderr}"
                logger.error(error_msg)
                return ""
        
        except subprocess.TimeoutExpired:
            error_msg = "cmake脚本执行超时"
            logger.error(error_msg)
            raise
        except Exception as e:
            error_msg = f"cmake执行脚本异常: {e}"
            logger.error(error_msg)
            raise

    def process_header_files(self, project_dir: str, header_dir: str, 
                           output_dir: str, taihec_script: str, max_retries: int = 3):
        """
        处理头文件工作流
        
        Args:
            project_dir: 项目目录
            header_dir: 头文件目录
            output_dir: 输出目录
            taihec_script: taihec 脚本路径
            max_retries: 最大重试次数
        """
        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)
        
        # 读取头文件
        header_files = self.read_header_files(header_dir)

        results = {}
        
        for filename, content in header_files.items():
            logger.info(f"开始处理文件: {filename}")
            #文件内容处理（忽略私有成员等）
            code=self.process_content(content,filename,project_dir)
            retry_count = 0
            success = False
            generated_code = ""
            error_history = []
            
            while retry_count <= max_retries and not success:
                try:
                    if retry_count == 0:
                        # 首次生成
                        generated_code = self.generate_code(code, filename)
                    else:
                        # 修复代码
                        generated_code = self.repair_code(
                            content, generated_code, 
                            "\n".join(error_history[-1:]),  # 最近3个错误
                            filename
                        )
                    
                    # 保存代码
                    code_file_path=self.save_generated_code(generated_code, output_dir, filename)
                    
                    # 运行检验脚本
                    success, error_msg = self.run_taihec_script(taihec_script, code_file_path,output_dir)
                    
                    #生成CMAKE文件
                    self.generate_cmakelist(taihec_script,output_dir,filename)
                    if not success:
                        error_history.append(f"尝试 {retry_count + 1}: {error_msg}")
                        retry_count += 1
                        logger.warning(f"{filename} 第 {retry_count} 次重试")
                    
                        
                except Exception as e:
                    error_msg = f"处理过程中异常: {e}"
                    error_history.append(error_msg)
                    retry_count += 1
                    logger.error(f"{filename} 处理异常: {e}")
            
            # 记录结果
            results[filename] = {
                'success': success,
                'retries': retry_count,
                'errors': error_history if not success else []
            }
            
            if success:
                logger.info(f"{filename} 处理成功")
            else:
                logger.error(f"{filename} 处理失败，已达到最大重试次数")
        
        # 输出总结报告
        self._generate_report(results, output_dir)
    
    def generate_cmakelist(self,script_path,output_dir: str,filename:str):
        lib_name = ""
        if filename.endswith('.d.ts'):
            lib_name  = filename[:-5]
        else:
            lib_name   = Path(filename).stem
        output_dir=Path(output_dir).as_posix()
        cmake =[]
        runtime_path = Path(self.search_taihe_dir(script_path,"--print-runtime-source-path")).as_posix()
        include_path = Path(self.search_taihe_dir(script_path,"--print-runtime-header-path")).as_posix()
        cmake.append("set(CMAKE_CXX_STANDARD 17)\n")
        cmake.append("set(CMAKE_CXX_STANDARD_REQUIRED True)\n\n")
        cmake.append("add_library(")
        cmake.append(f" {lib_name} SHARED\n")
        cmake.append(f" {output_dir}/temp/{lib_name}.impl.cpp\n")
        cmake.append(f" {output_dir}/temp/napi_register.cpp\n")
        cmake.append(f" {output_dir}/src/{lib_name}.napi.cpp\n")
        cmake.append(f" {runtime_path}/object.cpp\n")
        cmake.append(f" {runtime_path}/string.cpp\n")
        cmake.append(f" {runtime_path}/napi_runtime.cpp\n")
        cmake.append(")\n")

        cmake.append("target_include_directories(")
        cmake.append(f" {lib_name} PUBLIC\n")
        cmake.append(f" {include_path}\n")
        cmake.append(f" {output_dir}/include\n")
        cmake.append(")\n")
        with open(f"{output_dir}/CMakeLists.txt", "w", encoding="utf-8") as f:
            f.writelines(cmake)
            logger.info("cmake文件创建")


    def _generate_report(self, results: Dict, output_dir: str):
        """生成处理报告"""
        report_path = Path(output_dir) / "processing_report.txt"
        
        successful = sum(1 for r in results.values() if r['success'])
        failed = len(results) - successful
        
        with open(report_path, 'w', encoding='utf-8') as f:
            f.write("代码生成处理报告\n")
            f.write("=" * 50 + "\n\n")
            f.write(f"总文件数: {len(results)}\n")
            f.write(f"成功: {successful}\n")
            f.write(f"失败: {failed}\n\n")
            
            for filename, result in results.items():
                f.write(f"文件: {filename}\n")
                f.write(f"状态: {'成功' if result['success'] else '失败'}\n")
                f.write(f"重试次数: {result['retries']}\n")
                
                if not result['success']:
                    f.write("错误信息:\n")
                    for error in result['errors']:
                        f.write(f"  - {error}\n")
                f.write("\n")
        
        logger.info(f"处理报告已保存到: {report_path}")

# 使用示例


def load_config_from_json(config_path):
    """从JSON文件加载配置"""
    if not config_path.exists():
        raise FileNotFoundError(f"配置文件不存在: {config_path}")
    
    with open(config_path, 'r', encoding='utf-8') as f:
        return json.load(f)

def main():
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description="idl generator")
    parser.add_argument('-C', '--config', type=str, default='config.json', 
                  help='配置文件路径,默认为config.json')
    parser.add_argument('-H', '--header-dir', type=str, help='头文件目录,覆盖配置文件,需要转换的文件')
    parser.add_argument('-O', '--output-dir', type=str, help='输出目录,覆盖配置文件,默认是./generated')
    parser.add_argument('-P', '--project-dir', type=str, help='项目目录,覆盖配置文件,用于查询目标文件的头文件中的定义')
    parser.add_argument('-T', '--taihec-script', type=str, help='TAIHEC脚本路径,覆盖配置文件')
    parser.add_argument('-U', '--api-base-url', type=str, help='API基础URL,覆盖配置文件')
    parser.add_argument('-M', '--model-name', type=str, help='模型名称,覆盖配置文件')
    parser.add_argument('-R', '--mapping-rules-path', type=str, default='taihe_conversion_rules.json',
                  help='映射规则路径,覆盖配置文件')
    
    args = parser.parse_args()

    # 从环境变量获取API_KEY (优先于配置文件)
    API_KEY = os.getenv("API_KEY")
    if not API_KEY:
        raise ValueError("请设置环境变量 API_KEY")

    # 从JSON文件加载配置
    config_path = Path(args.config).resolve()
    config = load_config_from_json(config_path)

    # 解析配置，命令行参数 > JSON配置
    MAPPING_RULES_PATH = args.mapping_rules_path or config.get("mapping_rules_path")
    PROJECT_DIR = args.project_dir or config.get("project_dir")
    
    # 处理头文件目录
    header_dir_str = args.header_dir or config.get("header_dir")
    if not header_dir_str:
        raise ValueError("头文件目录未配置，请在配置文件或命令行中设置")
    HEADER_DIR = Path(header_dir_str).resolve()
    
    # 处理输出目录
    output_dir_str = args.output_dir or config.get("output_dir")
    if not output_dir_str:
        OUTPUT_DIR = (HEADER_DIR / "generated").resolve()
    else:
        OUTPUT_DIR = Path(output_dir_str).resolve()
    
    # 处理其他配置
    TAIHEC_SCRIPT = args.taihec_script or config.get("taihec_script")
    BASE_URL = args.api_base_url or config.get("api_base_url")
    MODEL_NAME = args.model_name or config.get("model_name")

    # 确保输出目录存在
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    
    # 初始化代码生成器
    generator = CodeGenerator(
        api_key=API_KEY,
        mapping_rules_path=MAPPING_RULES_PATH,
        base_url=BASE_URL,
        model_name=MODEL_NAME
    )
    
    # 执行处理流程
    generator.process_header_files(
        project_dir=PROJECT_DIR,
        header_dir=HEADER_DIR,
        output_dir=OUTPUT_DIR,
        taihec_script=TAIHEC_SCRIPT,
        max_retries=2
    )

if __name__ == "__main__":
    main()
